#include "onnxparser.h"
#include <fstream>
#include <algorithm>
#include "vkinfer/core/logger.h"

namespace vkinfer
{
	float OnnxAttribute::get_float(uint32_t index) const
	{
		return values.f32[index];
	}

	uint32_t OnnxAttribute::get_int(uint32_t index) const
	{
		return values.i32[index];
	}

	OnnxParser::OnnxParser() {}

	bool OnnxParser::open(const std::string& model_file)
	{
		std::ifstream input;
		input.open(model_file, std::ios::in | std::ios::binary);
		if (!input.is_open())
			return false;
		
		onnx::ModelProto model_proto;
		if (!model_proto.ParseFromIstream(&input))
			return false;

		if (!model_proto.has_graph())
			return false;

		graph = model_proto.graph();
		return true;
	}

	void OnnxParser::load_io_info(std::vector<std::string>& input_names, std::vector<std::string>& output_names)
	{
		input_names.clear();
		output_names.clear();

		for (int i = 0; i < graph.input_size(); ++i)
		{
			input_names.push_back(graph.input(i).name());
		}
		for (int i = 0; i < graph.output_size(); ++i)
		{
			output_names.push_back(graph.output(i).name());
		}
	}

	void OnnxParser::load_weights(std::map<std::string, std::shared_ptr<OnnxWeight>>& weights)
	{
		// weights are stored as initializers in onnx file
		weights.clear();
		weight_names.clear();

		for (int i = 0; i < graph.initializer_size(); ++i)
		{
			std::shared_ptr<OnnxWeight> weight = std::make_shared<OnnxWeight>();

			const onnx::TensorProto tensor = graph.initializer(i);
			weight->name = tensor.name();
			for (int k = 0; k < tensor.dims_size(); ++k)
			{
				uint32_t d = (uint32_t)tensor.dims(k);
				weight->shape.push_back(d);
			}
			// load weight values
			std::string raw_data_val = tensor.raw_data();
			uint32_t tensor_size = weight->count();
			const char* val = raw_data_val.c_str();
			float* values = new float[tensor_size];
			uint32_t index = 0;
			for (uint32_t k = 0; k <= tensor_size * 4 - 4; k += 4)
			{
				float w;
				char b[] = { val[k], val[k + 1], val[k + 2], val[k + 3] };
				memcpy(&w, &b, sizeof(float));
				values[index] = w;
				index += 1;
			}
			weight->values = values;
			weights.insert(std::make_pair(weight->name, weight));
			weight_names.insert(weight->name);
		}
	}

	void OnnxParser::load_network(std::map<std::string, std::shared_ptr<OnnxNodeInfo>>& nodes)
	{
		nodes.clear();

		for (int i = 0; i < graph.node_size(); ++i)
		{
			std::shared_ptr<OnnxNodeInfo> node_info = std::make_shared<OnnxNodeInfo>();
			const onnx::NodeProto node_proto = graph.node(i);
			
			node_info->name = node_proto.name();
#if 1
			if (node_info->name == "Constant_166")
			{
				std::cout << node_info->name << std::endl;
			}
#endif
			node_info->op_type = node_proto.op_type();
			
			// input/output
			for (int j = 0; j < node_proto.output_size(); ++j)
				node_info->output_names.push_back(node_proto.output(j));

			for (int j = 0; j < node_proto.input_size(); ++j)
			{
				std::string name = node_proto.input(j);

				if (weight_names.find(name) == weight_names.end())
					node_info->input_names.push_back(name);
				else
					node_info->weight_names.push_back(name);
			}

			// attributes
			for (int j = 0; j < node_proto.attribute_size(); ++j)
			{
				onnx::AttributeProto attrib_proto = node_proto.attribute(j);
				OnnxAttribute attribute = {};
				attribute.name = attrib_proto.name();
				auto type = attrib_proto.type();
				if (type == onnx::AttributeProto_AttributeType_FLOAT)
				{
					attribute.type = OnnxAttribValueType::VALUE_TYPE_FLOAT;
					attribute.num_values = 1;
					attribute.values.f32[0] = attrib_proto.f();
				}
				else if (type == onnx::AttributeProto_AttributeType_INT)
				{
					attribute.type = OnnxAttribValueType::VALUE_TYPE_INT;
					attribute.num_values = 1;
					attribute.values.i32[0] = (uint32_t)attrib_proto.i();
				}
				else if (type == onnx::AttributeProto_AttributeType_FLOATS)
				{
					assert(attrib_proto.floats_size() <= 4);
					attribute.type = OnnxAttribValueType::VALUE_TYPE_FLOATS;
					attribute.num_values = (uint32_t)attrib_proto.floats_size();
					for (int k = 0; k < attrib_proto.floats_size(); ++k)
					{
						if (k < 4) attribute.values.f32[k] = attrib_proto.floats(k);
					}
				}
				else if (type == onnx::AttributeProto_AttributeType_INTS)
				{
					attribute.type = OnnxAttribValueType::VALUE_TYPE_INTS;
					attribute.num_values = (uint32_t)attrib_proto.ints_size();
					for (int k = 0; k < attrib_proto.ints_size(); ++k)
						if (k < 4) attribute.values.i32[k] = (uint32_t)attrib_proto.ints(k);
				}
				else
				{
					attribute.num_values = 0;
				}
				node_info->attributes.push_back(attribute);
			}
			nodes.insert(std::make_pair(node_info->name, node_info));
		}
	}
}
